package main

import (
    "sync"
    "strings"
    "fmt"
    "errors"
    "encoding/binary"
    "io"
    "net"
)

type RdpProxy struct {
    proxy *TcpProxy
    clientAddr string
    serverAddr string
}

func newRdpProxy(conn *net.TCPConn) *RdpProxy {
    dp := &RdpProxy{}
    dp.proxy = newTcpProxy(conn)
    dp.clientAddr = conn.RemoteAddr().String()
    return dp
}

func (dp *RdpProxy)getAuth() (string, error) {
    buf := make([]byte, 4)
    _, err := io.ReadFull(dp.proxy.clientConn, buf)
    if err != nil {
        return "", err
    }
    
    authLen := binary.LittleEndian.Uint32(buf)
    if authLen <= 0 || authLen > 1024*1024 {
        return "", errors.New(fmt.Sprintf("invalid auth string len. (%d)", authLen))
    }
    buf = make([]byte, authLen)
    _, err = io.ReadFull(dp.proxy.clientConn, buf)
    if err != nil {
        return "", err
    }
    
    return string(buf), nil
}

func (dp *RdpProxy) runProxyAndWait() {
    defer dp.close()
    
    authStr, err := dp.getAuth()
    if err != nil {
        debugLog.Printf("get auth string failed. c(%s) err(%v)\n", dp.clientAddr, err)
        return
    }
    
    fields := strings.Split(authStr, ",")
    if len(fields) != 3 {
        debugLog.Printf("invalid auth string failed. c(%s) auth(%v)\n", dp.clientAddr, authStr)
        return
    }
    
    config := getConfig()
    // ip,port,password
    if config.RdpProxy.PassWord != "" && fields[2] != config.RdpProxy.PassWord {
        debugLog.Printf("invalid password string failed. c(%s) auth(%v)\n", dp.clientAddr, authStr)
        return
    }
    
    dp.serverAddr = fields[0] + ":" + fields[1]
    err = dp.proxy.connectToServer(dp.serverAddr)
    if err != nil {
        debugLog.Printf("connect to server failed. c(%s) s(%s) err(%v)\n", dp.clientAddr, dp.serverAddr, err)
        return
    }
    
    debugLog.Printf("start to proxy. c(%s) s(%s)\n", dp.clientAddr, dp.serverAddr)
    dp.proxy.runProxyAndWait()
}

func (dp *RdpProxy) close() {
    dp.proxy.close(nil)
}


type RdpProxyService struct {
    config  *RdpProxyConfig
    server  *TcpServer
    proxy   map[string]*RdpProxy
    lock    sync.Mutex
}

func newRdpProxyService(config *RdpProxyConfig) *RdpProxyService {
    s := &RdpProxyService{}
    s.config = config
    s.server = newTcpServer(config.ListenAddr)
    s.server.handler = s
    s.proxy = make(map[string]*RdpProxy)
    return s
}

func (s *RdpProxyService) service(conn *net.TCPConn) {
    dp := newRdpProxy(conn)
    success := s.addProxy(dp)
    if !success {
        dp.close()
        return
    }
    dp.runProxyAndWait()
    s.delProxy(dp)
}

func (s *RdpProxyService) addProxy(p *RdpProxy) bool {
    s.lock.Lock()
    defer s.lock.Unlock()
    
    s.proxy[p.clientAddr] = p
    
    debugLog.Printf("new rdp proxy. c(%s) total(%d)\n", p.clientAddr, len(s.proxy))
    return true
}

func (s *RdpProxyService) delProxy(p *RdpProxy) {
    s.lock.Lock()
    defer s.lock.Unlock()
    
    s.proxy[p.clientAddr] = nil
    delete(s.proxy, p.clientAddr)
    
    debugLog.Printf("del rdp proxy. c(%s) total(%d) err(%v)\n", p.clientAddr, len(s.proxy), p.proxy.err)
}

func (s *RdpProxyService) start() error {
    return s.server.start()
}

func (s *RdpProxyService) close() {
    s.lock.Lock()
    for _, p := range s.proxy {
        p.close()
    }
    s.lock.Unlock()
    
    s.server.close()
}

var gRdpProxyService *RdpProxyService = nil
func startRdpProxyService() {
    c := getConfig().RdpProxy
    gRdpProxyService = newRdpProxyService(c)
    err := gRdpProxyService.start()
    if err != nil {
        mainLog.Printf("start rdp proxy service failed. err(%v)\n", err)
    }
}

func reloadRdpProxyService() {
    c := getConfig().RdpProxy
    if c.ListenAddr == gRdpProxyService.config.ListenAddr {
        return
    }
    
    old := gRdpProxyService
    startRdpProxyService()
    old.close()
}